from tree_sitter import Language, Parser, Node
from entity import If_node, TRUE_TO_FALSE
from util import merge_var_change, format_compare
from z3 import Real, Solver, sat, And, Or, Not, simplify, Int
import re
from copy import deepcopy
from typing import Dict, List
from graphviz import Digraph
import os


def format_java_tree(root: Node, if_tree: If_node, vars: Dict[str, str]):
    """
    递归地处理 Java 代码的 AST 树，并构建 If_node 树结构。

    Parameters:
    - root: Tree-sitter 解析树的根节点。
    - if_tree: If_node 实例，表示当前的 if 树结构。
    - vars: 变量映射字典，用于更新和维护条件表达式中的变量值。
    """

    for child in root.children:
        # print(child,child.children, 1)
        if child.type == 'expression_statement':
            parts = child.children[0].text.decode('utf-8').split('=')
            c=child.children[0]
            # print(c.children)
            # print(parts)
            if parts.__len__() == 2:
                l_value = parts[0].strip()
                r_value = parts[1].strip()

                if(c.children[2].type == 'assignment_expression'):
                    vars[l_value] = l_value
                else:
                    if l_value in vars:
                        vars[l_value] = vars[l_value].replace(l_value, f'({r_value})')
                    else:
                        vars[l_value] = r_value
        # print(child,1)
        if child.type == 'if_statement':
            cc = child.children
            # print(cc,2)
            # 处理 if 语句，并将条件表达式添加到 if_node 树中
            condition = merge_var_change(str(cc[1].text, encoding='utf-8').strip(), vars)
            if_tree.children.append(If_node('if', condition))
            # for child in cc[2].children[1:]:
            #     print(child,3)
            #     print(child.children,4)
            format_java_tree(cc[2], if_tree.children[-1], vars)
            if cc.__len__()>3:
                if cc[3].type == 'else':
                    if_tree.children.append(If_node('else', None))
                    format_java_tree(cc[4], if_tree.children[-1], deepcopy(vars))
                if cc[3].type == 'else_if_statement':
                    child=cc[3].children
                    condition = merge_var_change(str(child[1].text, encoding='utf-8').strip(), vars)
                    if_tree.children.append(If_node('else_if_statement', condition))
                    format_java_tree(child[2], if_tree.children[-1], vars)
            # 处理 else if 和 else 分支
            # for c in cc[2:]:
            #     print(c,214231)
            #     if c.type == 'block':
            #         format_java_tree(c.children[1], if_tree.children[-1], deepcopy(vars))
            #     if c.type == 'else':
            #         if_tree.children.append(If_node('else', None))
            #     elif c.type == 'else_if_statement':
            #         condition = merge_var_change(str(c.children[1].text, encoding='utf-8').strip(), vars)
            #         if_tree.children.append(If_node('else_if_clause', condition))
            #         format_java_tree(c.children[3], if_tree.children[-1], deepcopy(vars))
            #     for child in c.children:
            #         if child.type == 'if_statement':
            #             cc = child.children
            #             print(cc,cc.__len__(),cc[0].text,cc[1].text,"nbbb")
            #             # 处理 if 语句，并将条件表达式添加到 if_node 树中
            #             condition = merge_var_change(str(cc[1].text, encoding='utf-8').strip(), vars)
            #             if_tree.children.append(If_node('if_statement', condition))
            #             if cc.__len__()>3:
            #                 format_java_tree(cc[3], if_tree.children[-1], deepcopy(vars))
            #         if child.type == 'else':
            #             if_tree.children.append(If_node('else', None))
            #         elif child.type == 'else_if_statement':
            #             condition = merge_var_change(str(c.children[1].text, encoding='utf-8').strip(), vars)
            #             if_tree.children.append(If_node('else_if_clause', condition))
            #             format_java_tree(c.children[3], if_tree.children[-1], deepcopy(vars))
            #         if child.type == 'else_if_clause':
            #             condition = merge_var_change(str(c.children[1].text, encoding='utf-8').strip(), vars)
            #             if_tree.children.append(If_node('else_if_clause', condition))
            #             format_java_tree(c.children[3], if_tree.children[-1], deepcopy(vars))
            #         elif child.type == 'else':
            #             if_tree.children.append(If_node('else', None))
            #             format_java_tree(c.children[1], if_tree.children[-1], deepcopy(vars))


def merge_var_change(expression: str, vars: Dict[str, str]) -> str:
    """
    替换表达式中的变量为其对应的值。

    Parameters:
    - expression: 原始表达式。
    - vars: 变量映射字典。

    Returns:
    - 替换后的表达式。
    """
    for var, value in vars.items():
        expression = re.sub(r'\b{}\b'.format(re.escape(var)), value, expression)
    return expression


def get_java_truth_trees(root: If_node) -> List[List[str]]:
    def get_re(root: If_node, prefix: list, res: List[List[str]]):
        if root.compare_statement != None:
            # 格式化条件并加入到prefix
            prefix.append(format_compare(root.compare_statement))
        i = root.children.__len__()
        # 当前节点没有子节点
        if i == 0:
            if prefix != []:
                # 将prefix加入到res
                res.append(prefix)
            return
        # 遍历当前节点的所有子节点
        for ii in range(i):
            # print(root.children[ii], deepcopy(prefix),124)
            # 从当前节点的第ii个子节点递归调用函数，生成从该节点开始的路径约束
            # deepcopy(prefix)是对prefix深拷贝
            get_re(root.children[ii], deepcopy(prefix), res)
            if root.children[ii].compare_statement != None:
                # print(222)
                # 将当前子节点的条件的否定添加到prefix
                prefix.append(f'Not({format_compare(root.children[ii].compare_statement)})')

    l = len(root.children)
    res: List[List[str]] = []
    prefix = []  # 到达这个分支需要的前置条件
    for i in range(l):
        get_re(root.children[i], deepcopy(prefix), res)
        # print(root.children[i].compare_statement,root.children[i],res)
        if root.children[i].compare_statement != None:
            # print(222)
            prefix.append(f'Not({format_compare(root.children[i].compare_statement)})')
    return res
    # 如果当前比较条件不是none

def print_java_tree(root: If_node, level=0):
    print('  ' * level + repr(root))
    for child in root.children:
        print_java_tree(child, level + 1)


def replace_java_keys_with_values(string: str, rep_dict: dict):
    if string is None:
        return None
    for key, value in rep_dict.items():
        if key in string:
            string = string.replace(key, value)
    return string


def get_java_variables(s: str):
    variables = re.findall(r"\b[a-zA-Z]+\b", s)
    variables = list(filter(lambda x: x not in ['And', 'Not', 'Or'], variables))
    return variables


JAVA_LANGUAGE = Language('build/my-languages.so', 'java')
java_parser = Parser()
java_parser.set_language(JAVA_LANGUAGE)


def analyzeJava(code: str):
    result = ''
    tree = java_parser.parse(bytes(code, "utf8"))
    root_node = tree.root_node
    if_tree = If_node('module', None)
    format_java_tree(root_node, if_tree, {})
    restrictions = get_java_truth_trees(if_tree)

    letters = set()
    for res in restrictions:
        for r in res:
            v = get_java_variables(r)
            letters.update(v)

    variables = {'And': And, 'Or': Or, 'Not': Not}
    for x in letters:
        variables[x] = Int(x)

    result += '一组路径覆盖测试用例如下\n'

    for res in restrictions:
        s = Solver()
        for r in res:
            rr = simplify(eval(r, {}, variables))
            s.add(rr)
        result += f'该测试用例约束条件为:{res}\n'
        if s.check() == sat:
            m = s.model()
            result += str(m) + '\n'
        else:
            result += "无解\n"
    return result



java_code = """
a=a+1;
if(a>2){
    System.out.println(999);
    if(a>3){
        System.out.println(a);
    }
    else{
        System.out.println(a);
    }
}
else{
    System.out.println(a);
}
"""
# print(analyzeJava(java_code))
